import os
import re

import tqdm  # 导入 tqdm 以显示处理进度
from datasets import load_dataset

# 定义数据集保存路径
DATASET_DIR = "bert_pytorch/dataset"
os.makedirs(DATASET_DIR, exist_ok=True)

# 定义保存文件的路径
train_file_path = os.path.join(DATASET_DIR, "wikitext-2-train.txt")
valid_file_path = os.path.join(DATASET_DIR, "wikitext-2-valid.txt")
test_file_path = os.path.join(DATASET_DIR, "wikitext-2-test.txt")

# 加载 wikitext-2 数据集
wikitext_dataset = load_dataset("wikitext", "wikitext-2-raw-v1")


def clean_text(text):
    """
    对单行文本数据进行清洗，包括去除标题行、标准化特殊标记、转换为小写、去除多余空格。

    Args:
        text: 原始文本字符串

    Returns:
        cleaned_text: 清洗后的文本字符串
    """
    cleaned_text = text.strip()  # 去除行首尾空白

    # 1. 去除标题/章节行 (以 '= =' 开头和结尾的行)
    if re.match(r'^= =.*?= =$', cleaned_text):
        return None  # 返回 None 表示此行应被忽略

    if re.match(r'^=.*?=$', cleaned_text):
        return None  # 返回 None 表示此行应被忽略

    if re.match(r'^[^\w\s].*[^\w\s]$', cleaned_text):
        return None  # 返回 None 表示此行应被忽略
    # 2. 标准化特殊标记 '@-@' 和 '@,@'
    cleaned_text = cleaned_text.replace(' @-@ ', '-')  # 注意前后加空格，避免误替换
    cleaned_text = cleaned_text.replace(' @,@ ', ',')  # 注意前后加空格，避免误替换

    # 3. 转换为小写 (可选，这里选择转换)
    cleaned_text = cleaned_text.lower()

    # 4. 处理多余空格 (使用 split() 和 join() 方法)
    cleaned_text = ' '.join(cleaned_text.split())

    return cleaned_text


def save_as_sentence_pairs(dataset, file_path):
    """
    处理数据集，将相邻两句组合成句对格式，并保存到文件。
    在保存前，对每行文本进行数据清洗。
    """
    print(f"开始处理数据集并保存到: {file_path}")
    with open(file_path, 'w', encoding='utf-8') as f:
        lines = []
        for example in tqdm.tqdm(dataset):  # 使用 tqdm 显示处理进度
            original_text = example['text'].strip()
            if not original_text:  # 跳过空行
                continue

            cleaned_line = clean_text(original_text)  # 对每一行进行清洗
            if cleaned_line:  # 确保清洗后不是 None (例如，标题行清洗后返回 None)
                lines.append(cleaned_line)

        sentence_pairs_count = 0
        for i in range(len(lines) - 1):
            f.write(lines[i] + "\t" + lines[i + 1] + "\n")  # 用 \t 连接两句话
            sentence_pairs_count += 1
        print(f"数据清洗和句对生成完成，共保存 {sentence_pairs_count} 个句对到: {file_path}")


# 处理并保存数据 (清洗过程已包含在 save_as_sentence_pairs 函数中)
print("开始处理训练集...")
save_as_sentence_pairs(wikitext_dataset['train'], train_file_path)
print("\n开始处理验证集...")
save_as_sentence_pairs(wikitext_dataset['validation'], valid_file_path)
print("\n开始处理测试集...")
save_as_sentence_pairs(wikitext_dataset['test'], test_file_path)

print("\n所有数据集处理完成！")
